"""
存储管理器
管理数据文件的创建、删除和访问
"""

import os
import struct
import json
from typing import Dict, List, Optional, Any

class StorageManager:
    """存储管理器"""
    
    def __init__(self, data_dir: str = "data"):
        """初始化存储管理器"""
        self.data_dir = data_dir
        self.page_size = 4096  # 4KB页面大小
        self.meta_file = os.path.join(data_dir, "metadata.json")
        self.tables = {}  # 表信息缓存
        
        # 确保数据目录存在
        os.makedirs(data_dir, exist_ok=True)
        
        # 加载元数据
        self.load_metadata()
    
    def load_metadata(self):
        """加载元数据"""
        if os.path.exists(self.meta_file):
            try:
                with open(self.meta_file, 'r', encoding='utf-8') as f:
                    self.tables = json.load(f)
            except Exception as e:
                print(f"加载元数据失败: {e}")
                self.tables = {}
    
    def save_metadata(self):
        """保存元数据"""
        try:
            with open(self.meta_file, 'w', encoding='utf-8') as f:
                json.dump(self.tables, f, indent=2, ensure_ascii=False)
        except Exception as e:
            print(f"保存元数据失败: {e}")
    
    def create_table(self, table_name: str, schema: Dict[str, Any]) -> bool:
        """创建表"""
        try:
            # 检查表是否已存在
            if table_name in self.tables:
                raise Exception(f"表 '{table_name}' 已存在")
            
            # 创建表文件
            table_file = os.path.join(self.data_dir, f"{table_name}.dat")
            
            # 初始化表信息
            table_info = {
                'name': table_name,
                'file': table_file,
                'schema': schema,
                'page_count': 0,
                'record_count': 0,
                'next_page_id': 0,
                'free_pages': []
            }
            
            # 创建初始页面
            self._create_initial_page(table_file)
            table_info['page_count'] = 1
            table_info['next_page_id'] = 1
            
            # 保存表信息
            self.tables[table_name] = table_info
            self.save_metadata()
            
            return True
            
        except Exception as e:
            print(f"创建表失败: {e}")
            return False
    
    def drop_table(self, table_name: str) -> bool:
        """删除表"""
        try:
            if table_name not in self.tables:
                raise Exception(f"表 '{table_name}' 不存在")
            
            # 删除表文件
            table_file = self.tables[table_name]['file']
            if os.path.exists(table_file):
                os.remove(table_file)
            
            # 从元数据中删除
            del self.tables[table_name]
            self.save_metadata()
            
            return True
            
        except Exception as e:
            print(f"删除表失败: {e}")
            return False
    
    def get_table_info(self, table_name: str) -> Optional[Dict[str, Any]]:
        """获取表信息"""
        return self.tables.get(table_name)
    
    def list_tables(self) -> List[str]:
        """列出所有表"""
        return list(self.tables.keys())
    
    def _create_initial_page(self, table_file: str):
        """创建初始页面"""
        # 页面头部结构: [页面ID(4字节)][记录数(4字节)][空闲空间(4字节)][数据...]
        page_header_size = 12
        data_size = self.page_size - page_header_size
        
        # 创建页面数据
        page_data = bytearray(self.page_size)
        
        # 写入页面头部
        struct.pack_into('<I', page_data, 0, 0)  # 页面ID
        struct.pack_into('<I', page_data, 4, 0)  # 记录数
        struct.pack_into('<I', page_data, 8, data_size)  # 空闲空间
        
        # 写入文件
        with open(table_file, 'wb') as f:
            f.write(page_data)
    
    def allocate_page(self, table_name: str) -> int:
        """分配新页面"""
        if table_name not in self.tables:
            raise Exception(f"表 '{table_name}' 不存在")
        
        table_info = self.tables[table_name]
        
        # 检查是否有空闲页面
        if table_info['free_pages']:
            page_id = table_info['free_pages'].pop()
        else:
            # 创建新页面
            page_id = table_info['next_page_id']
            table_info['next_page_id'] += 1
            table_info['page_count'] += 1
            
            # 扩展文件
            table_file = table_info['file']
            with open(table_file, 'ab') as f:
                f.write(b'\x00' * self.page_size)
        
        self.save_metadata()
        return page_id
    
    def deallocate_page(self, table_name: str, page_id: int):
        """释放页面"""
        if table_name not in self.tables:
            raise Exception(f"表 '{table_name}' 不存在")
        
        table_info = self.tables[table_name]
        table_info['free_pages'].append(page_id)
        self.save_metadata()
    
    def read_page(self, table_name: str, page_id: int) -> Optional[bytearray]:
        """读取页面"""
        if table_name not in self.tables:
            raise Exception(f"表 '{table_name}' 不存在")
        
        table_file = self.tables[table_name]['file']
        
        try:
            with open(table_file, 'rb') as f:
                f.seek(page_id * self.page_size)
                page_data = f.read(self.page_size)
                if len(page_data) == self.page_size:
                    return bytearray(page_data)
                else:
                    return None
        except Exception as e:
            print(f"读取页面失败: {e}")
            return None
    
    def write_page(self, table_name: str, page_id: int, page_data: bytearray) -> bool:
        """写入页面"""
        if table_name not in self.tables:
            raise Exception(f"表 '{table_name}' 不存在")
        
        if len(page_data) != self.page_size:
            raise Exception(f"页面大小不正确: {len(page_data)} != {self.page_size}")
        
        table_file = self.tables[table_name]['file']
        
        try:
            with open(table_file, 'r+b') as f:
                f.seek(page_id * self.page_size)
                f.write(page_data)
            return True
        except Exception as e:
            print(f"写入页面失败: {e}")
            return False
    
    def get_page_count(self, table_name: str) -> int:
        """获取表的页面数"""
        if table_name not in self.tables:
            return 0
        return self.tables[table_name]['page_count']
    
    def get_record_count(self, table_name: str) -> int:
        """获取表的记录数"""
        if table_name not in self.tables:
            return 0
        return self.tables[table_name]['record_count']
    
    def update_record_count(self, table_name: str, delta: int):
        """更新记录数"""
        if table_name in self.tables:
            self.tables[table_name]['record_count'] += delta
            self.save_metadata()
    
    def get_schema(self, table_name: str) -> Optional[Dict[str, Any]]:
        """获取表结构"""
        if table_name not in self.tables:
            return None
        return self.tables[table_name]['schema']
    
    def table_exists(self, table_name: str) -> bool:
        """检查表是否存在"""
        return table_name in self.tables
    
    def query_table(self, table_name: str, conditions: Dict[str, Any]) -> List[Dict[str, Any]]:
        """查询表数据（简化实现）"""
        if not self.table_exists(table_name):
            return []
        
        # 简化实现：从内存中查询
        # 这里应该实现实际的查询逻辑，但为了简化，我们使用内存存储
        if not hasattr(self, '_memory_data'):
            self._memory_data = {}
        
        if table_name not in self._memory_data:
            self._memory_data[table_name] = []
        
        # 简单的条件匹配
        results = []
        for record in self._memory_data[table_name]:
            match = True
            for key, value in conditions.items():
                if record.get(key) != value:
                    match = False
                    break
            if match:
                results.append(record)
        
        return results
    
    def insert_record(self, table_name: str, record: Dict[str, Any]) -> bool:
        """插入记录（简化实现）"""
        if not self.table_exists(table_name):
            return False
        
        # 简化实现：存储到内存
        if not hasattr(self, '_memory_data'):
            self._memory_data = {}
        
        if table_name not in self._memory_data:
            self._memory_data[table_name] = []
        
        self._memory_data[table_name].append(record)
        return True
    
    def update_record(self, table_name: str, conditions: Dict[str, Any], updates: Dict[str, Any]) -> bool:
        """更新记录（简化实现）"""
        if not self.table_exists(table_name):
            return False
        
        # 简化实现：更新内存中的数据
        if not hasattr(self, '_memory_data') or table_name not in self._memory_data:
            return False
        
        updated = False
        for record in self._memory_data[table_name]:
            match = True
            for key, value in conditions.items():
                if record.get(key) != value:
                    match = False
                    break
            if match:
                record.update(updates)
                updated = True
        
        return updated
    
    def delete_record(self, table_name: str, conditions: Dict[str, Any]) -> bool:
        """删除记录（简化实现）"""
        if not self.table_exists(table_name):
            return False
        
        # 简化实现：从内存中删除数据
        if not hasattr(self, '_memory_data') or table_name not in self._memory_data:
            return False
        
        original_count = len(self._memory_data[table_name])
        self._memory_data[table_name] = [
            record for record in self._memory_data[table_name]
            if not all(record.get(key) == value for key, value in conditions.items())
        ]
        
        return len(self._memory_data[table_name]) < original_count
    
    def print_storage_info(self):
        """打印存储信息"""
        print("存储信息:")
        print(f"  数据目录: {self.data_dir}")
        print(f"  页面大小: {self.page_size} 字节")
        print(f"  表数量: {len(self.tables)}")
        
        for table_name, table_info in self.tables.items():
            print(f"  表: {table_name}")
            print(f"    文件: {table_info['file']}")
            print(f"    页面数: {table_info['page_count']}")
            print(f"    记录数: {table_info['record_count']}")
            print(f"    空闲页面: {len(table_info['free_pages'])}")

# 测试函数
if __name__ == "__main__":
    storage = StorageManager("test_data")
    
    # 测试创建表
    schema = {
        'columns': {
            'id': {'type': {'type': 'int'}, 'primary_key': True},
            'name': {'type': {'type': 'varchar', 'length': 50}},
            'age': {'type': {'type': 'int'}}
        },
        'primary_key': 'id'
    }
    
    print("创建表...")
    if storage.create_table("users", schema):
        print("表创建成功")
    
    print("\n分配页面...")
    page_id = storage.allocate_page("users")
    print(f"分配页面ID: {page_id}")
    
    print("\n存储信息:")
    storage.print_storage_info()
    
    print("\n删除表...")
    if storage.drop_table("users"):
        print("表删除成功")
